import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, explained_variance_score

def calculate_metrics(y_true, y_pred):
    """增加更多评估指标"""
    metrics = {
        'MSE': mean_squared_error(y_true, y_pred),
        'RMSE': np.sqrt(mean_squared_error(y_true, y_pred)),
        'MAE': mean_absolute_error(y_true, y_pred),
        'R2': r2_score(y_true, y_pred),
        'MAPE': np.mean(np.abs((y_true - y_pred) / (y_true + 1e-8))) * 100,
        'Explained_Variance': explained_variance_score(y_true, y_pred)
    }
    return metrics

def print_metrics(metrics):
    """打印评估指标"""
    for name, value in metrics.items():
        print(f'{name}: {value:.4f}') 

def evaluate_prediction_quality(metrics):
    """评估预测质量"""
    quality_assessment = {
        'excellent': False,
        'good': False,
        'acceptable': False,
        'needs_improvement': True
    }
    
    if metrics['R2'] > 0.9 and metrics['MAPE'] < 10:
        quality_assessment['excellent'] = True
        quality_assessment['needs_improvement'] = False
    elif metrics['R2'] > 0.8 and metrics['MAPE'] < 20:
        quality_assessment['good'] = True
        quality_assessment['needs_improvement'] = False
    elif metrics['R2'] > 0.7 and metrics['MAPE'] < 30:
        quality_assessment['acceptable'] = True
        quality_assessment['needs_improvement'] = False
    
    return quality_assessment 